package com.bsj.media.media;

import com.bsj.power.common.def.protocol.rtp.RtpParameterSet;
import com.bsj.power.common.util.Conv;
import lombok.extern.slf4j.Slf4j;


import java.util.Arrays;

/**
 * @ClassName ParserH265
 * @Description 解析H265数据
 * @Author zhoushubin
 * @Date 2022/3/24 15:36
 **/
@Slf4j
public class ParserH265 {


    static final int RET_OK = 0;
    static final int RET_ERR_FILE_NOT_EXIST = -1;
    static final int RET_ERR_UNKNOWN = -2;
    static final int RET_ERR_FILE_END = -3;
    static final int RET_ERR = -4;
    static final int RET_ERR_END_NAL_NOT_FOUND = -5;  // 表示没有下一帧
    static final int RET_ERR_START_CODE_NOT_FOUND = -6; // 当天帧错误

    public static int ReadOneNaluFromBuf(byte[] buffer, int nBufferSize, int offSet, NalInfo pNalu) {
        int startPos = 0, endPos = 0;
        int remain = 0;
        int nalFirstByte; // Nalu类型(从左到右开始算Xbit) 0bit:f  1bit-6bit:NaLType  7bit-12bit:LayerId  13bit-15bit:Tid
        int start0_code_start_size = 0;

        if (offSet + 3 >= nBufferSize) {
            return RET_ERR_START_CODE_NOT_FOUND;
        }

        startPos = FindStartCodeEx(buffer, offSet, nBufferSize - offSet);
        if (startPos < 0) {
            return RET_ERR_START_CODE_NOT_FOUND;
        }

        start0_code_start_size = (buffer[startPos + 2] == 0x1) ? 3 : 4;
        remain = nBufferSize - (startPos + start0_code_start_size);
        if (remain < 3) {
            return RET_ERR_END_NAL_NOT_FOUND;
        }

        endPos = FindStartCodeEx(buffer, startPos + start0_code_start_size, remain);
        if (endPos < 0) {
            return RET_ERR_END_NAL_NOT_FOUND;
        }

        pNalu.len = endPos - startPos; // 包括00 00 00 01
        pNalu.buf = Arrays.copyOfRange(buffer, startPos, endPos);
        nalFirstByte = (pNalu.buf[2] == 0x1) ? Conv.getShortNetOrder(pNalu.buf, 3) : Conv.getShortNetOrder(pNalu.buf, 4);
        pNalu.forbidden_bit = (nalFirstByte & 0x8000) > 0;
        //pNalu.nal_reference_idc = nalFirstByte & 0x60;
        pNalu.nal_unit_type = (nalFirstByte & 0x7E00) >> 9;
        return endPos;
    }

    /**
     * @Description: 从当前offSet读取一帧数据赋值到nalu
     * @param buffer: 原始数据
     * @param nBufferSize: 原始数据长度
     * @param offSet: 当前帧的起始偏移量
     * @param nalu: 存储一帧信息
     * @return: int 负数表示读取帧失败，正数==offSet(表示正常帧的起始帧)
     **/
    public static int realOneNalu(byte[] buffer, int nBufferSize, int offSet, NalInfo nalu) {
        int nalFirstByte; // Nalu类型(从左到右开始算Xbit) 0bit:f  1bit-6bit:NaLType  7bit-12bit:LayerId  13bit-15bit:Tid
        int start0_code_start_size = 0;
        int startPos = 0, endPos = 0;
        int remain = 0;
        if (offSet + 3 >= nBufferSize) {
            return -1;
        }

        startPos = FindStartCodeEx(buffer, offSet, nBufferSize - offSet);
        if (startPos < 0) {
            return -1;
        }

        start0_code_start_size = (buffer[startPos + 2] == 0x1) ? 3 : 4;
        remain = nBufferSize - (startPos + start0_code_start_size);// 下帧剩余长度
        if (remain < 3) {
            return -1;
        }

        endPos = FindStartCodeEx(buffer, startPos + start0_code_start_size, remain);
        if (endPos < 0) {
            // 无下帧
            endPos = nBufferSize; // buffer最后的偏移量 + 1
        }
        nalu.len = endPos - startPos; // 包括00 00 00 01
        nalu.buf = Arrays.copyOfRange(buffer, startPos, endPos);
        nalFirstByte = (nalu.buf[2] == 0x1) ? Conv.getShortNetOrder(nalu.buf, 3) : Conv.getShortNetOrder(nalu.buf, 4);
        nalu.forbidden_bit = (nalFirstByte & 0x8000) > 0;
        //pNalu.nal_reference_idc = nalFirstByte & 0x60;
        nalu.nal_unit_type = (nalFirstByte & 0x7E00) >> 9;
        return offSet;
    }

    public static int FindStartCodeEx(byte[] buf, int offset, int size) {
        boolean bFound = false;
        int pos = offset;
        while (pos + 4 <= size && pos < buf.length) {
            if (buf[pos] != 0x0 || buf[pos + 1] != 0x0) {
                pos++;
                continue;
            }

            bFound = (buf[pos + 2] == 0x1) || (buf[pos + 2] == 0x0 && buf[pos + 3] == 0x1);
            if (bFound) {
                return pos;
            }
            pos += 1;
        }

        return -1;
    }

    /**
     * @Description: 获取I帧中的vps、sps以及pps数据
     * @param body:  H265数据
     * @param srcOffset: H265起始位置偏移值
     * @param vps:
     * @param sps:
     * @param pps:
     * @return: int 返回关键帧的起始位置，去掉了H265帧头
     **/
    public static int readIFrameVPSSPSPPSFromBuf(byte[] body, int srcOffset, NalInfo vps, NalInfo sps, NalInfo pps) {
        int mediaLen = 0;
        NalInfo nalu = new NalInfo();

        // 读取数据vps、sps、pps
        int ret = realOneNalu(body, body.length, srcOffset, vps);
        if (ret < 0 || vps.getNal_unit_type() != RtpParameterSet.naluTypeVPSH5) {
            return -32; // 随便返回
        }
        srcOffset = ret + vps.len;

        ret = realOneNalu(body, body.length, srcOffset, sps);
        if (ret < 0 || sps.getNal_unit_type() != RtpParameterSet.naluTypeSPSH5) {
            return -33;
        }
        srcOffset = ret + sps.len;

        ret = realOneNalu(body, body.length, srcOffset, pps);
        if (ret < 0 || pps.getNal_unit_type() != RtpParameterSet.naluTypePPSH5) {
            return -34;
        }
        srcOffset = ret + pps.len;

        // 读取关键帧数据，中间有sei帧，需要判断和过滤掉
        ret = realOneNalu(body, body.length, srcOffset, nalu);

        while (ret > 0 && nalu.nal_unit_type == RtpParameterSet.naluTypeSEIH5) {
            srcOffset = ret + nalu.len;
            nalu = new NalInfo();
            ret = realOneNalu(body, body.length, srcOffset, nalu);
        }

        if (ret < 0) {
            mediaLen = ret;
            log.info("H265异常数据，未解析到帧数据");
        } else if (nalu.nal_unit_type == RtpParameterSet.naluTypeIDRH5) {
            // 查询到关键帧
            mediaLen = ret + ((nalu.buf[2] == 0x1) ? 3 : 4);
        } else {
            // 未查询到关键帧 理论上不可能到这里
            log.info("未查询到关键帧数据：ret:" + ret + ", type:" + nalu.getNal_unit_type());
        }

        return mediaLen;
    }
}
